﻿from typing import Any
import numpy as np
import scipy.io as scio
import os
import torch


class LoadBCICompDataSet:
    '''
    这个类提供一些简单的函数去加载BCI Comp III的内容

    刺激方法：
        数字化为240Hz
        6行6列，行闪和列闪，重复15次
        矩阵会展示2.5秒，这期间所有的字符亮度一致，所以这个数据最长也就是2.5*240
        之后，一行或一列亮100ms
        每次闪烁后，矩阵会空白75ms，每行每列重复15次
        一个循环结束后，显示2.5s的空白

    数据集内容：
        训练共有85个字符，测试共有100字符
        分为A和B两个人
        A TargetChar EAEVQTDOJG8RBRGONCEDHCTUIDBPUHMEM6OUXOCFOUKWA4VJEFRZROLHYNQDW_EKTLBWXEPOUIKZERYOOTHQI
        B TargetChar VGREAAH8TVRHBYN_UGCOLO4EUERDOOHCIFOMDNU6LQCPKEIREKOYRQIDJXPBKOJDWZEUEWWFOEBHXTQTTZUMO
    
    Mat文件内容及格式：
    Flashing        single  85x7794     是否是激发变亮的
    Signal          single  85x7794x64
    StimulusCode    single  85x7794     0是没有列/行激发，1-6是激发的列，7-12是激发的行
    StimulusType    single  85x7794     0是没有列/行激发，1是激发的行/列包含需要的字符，理论上有2550个，实际上根据示例代码判断的画也有2550个，但是signal有64个通道，所以是163200个数据
    TargetChar      char    1×85

    PDF给出的提取方法
        提取与特定行/列加强相关的信号波形仅需几个步骤：
            • 对于一个或多个通道，收集每次加强开始时的信号样本时间段，即每当闪烁从0变为1时（注意：数据集中的每个字符周期都是从第一次闪烁开始，即每个周期的第一个数据样本中闪烁=1）。
            • 根据相应刺激的刺激代码，将信号样本累加到12个单独的缓冲区中。对于每个字符周期，每个缓冲区应包含15个样本周期 - 对应给定行/列的每次加强

    当Flash为0的时候，stimulateType也为0

    刺激用的字符
        1   2   3   4   5   6                   0   1   2   3   4   5
    7   A   B   C   D   E   F               6   A   B   C   D   E   F
    8   G   H   I   J   K   L               7   G   H   I   J   K   L
    9   M   N   O   P   Q   R      ->       8   M   N   O   P   Q   R
    10  S   T   U   V   W   X               9   S   T   U   V   W   X
    11  Y   Z   1   2   3   4               10  Y   Z   1   2   3   4
    12  5   6   7   8   9   _               11  5   6   7   8   9   _

    左边一列和上面一行的数字是是数据集使用的索引
    
    Attribute:
        filePath：文件所在的路径
        _dict：访问类内的成员变量
        flashing：存储文件的Flashing变量的内容
        signal：存储文件的Signal变量的内容
        sitmulusCode：存储文件的SitmulusCode变量的内容
        sitmulusType：存储文件的SitmulusType变量的内容
        TargetChar：存储文件的TargetChar变量的内容
        colAddRowNum：刺激界面的col加上row的大小
        colMulRowNum：刺激界面的col乘上row的大小
        times：实验的次数
        totalDataNum：总共的时间点
        windows：数据需处理的长度，因为该数据集的采样率是240Hz，所以240是1s
        colAndRowApper：行和列出现的次数，行和列是一样的
        responses：处理后的每次实验的所在行或列的刺激的信号信息
        targetRowAndCol：目标字符所在的行和列数组
        targetResponses：目标字符对应的信号信息
        allLetterResponse：所有字符对应的刺激信号
        targetLetterAvgResponses：字符对应的多次刺激信息的平均
        targetLetterResponses：字符对应的多次刺激信息
        ScreenLetter：屏幕的字符数组
    '''

    def __init__(self, filePath: str,
                 col: int = 6,
                 row: int = 6,
                 times: int = 85,
                 totalDataNum: int = 7794,
                 channels: int = 64,
                 windows: int = 240,
                 colAndRowAppear: int = 15) -> None:
        '''
        Description:
            filePath： mat文件路径
            col: 刺激范式界面的列数
            row: 刺激范式界面的行数
            times: 实验重复的次数
            totalDataNum: 记录数据的总数
            channels: 数据的通道数
            windows: 数据需处理的长度，因为该数据集的采样率是240Hz，所以240是1s
            colAndRpwAppear: 一次实验重复的次数
        '''
        self.filePath = filePath
        self._dict = None
        self.data = None
        self.flashing = None
        self.signal = None
        self.stimulusCode = None
        self.stimulusType = None
        self.targetChar = None
        self.colAddRowNum = col + row
        self.colMulRowNum = row * col
        self.times = times
        self.totalDataNum = totalDataNum
        self.channels = channels
        self.windows = windows
        self.colAndRowAppear = colAndRowAppear
        self.responses = np.zeros(
            (self.times, self.colAddRowNum, self.colAndRowAppear, self.windows, self.channels))  # 15是每个信号的行和列出现15次
        self.targetRowAndCol = np.ndarray((85, 3))
        self.targetResponses = np.zeros((self.times, 2, self.colAndRowAppear, self.windows, self.channels))
        self.allLetterResponse = np.zeros((self.times, self.colMulRowNum, self.windows, self.channels))
        self.targetLetterAvgResponses = np.zeros((self.times, self.windows, self.channels))
        self.targetLetterResponses = np.zeros((self.times, self.colAndRowAppear, self.windows, self.channels))
        self.ScreenLetter = ["A", "B", "C", "D", "E", "F",
                             "G", "H", "I", "J", "K", "L",
                             "M", "N", "O", "P", "Q", "R",
                             "S", "T", "U", "V", "W", "X",
                             "Y", "Z", "1", "2", "3", "4",
                             "5", "6", "7", "8", "9", "_"]
        if os.path.exists(filePath):
            self.LoadData()
            self.__makeDict()

    def __makeDict(self) -> None:
        '''
        Description:
            制作变量字典
            通过字典和__getitem__、__setitem__让使用者以'[]'的方式获取和设置变量值
        '''
        self._dict = {
            "Flashing": self.flashing,
            "Signal": self.signal,
            "StimulusCode": self.stimulusCode,
            "StimulusType": self.stimulusType,
            "TargetChar": self.targetChar,
            "Responses": self.responses,
            "TargetRowAndCol": self.targetRowAndCol,
            "TargetResponses": self.targetResponses,
            "AllLetterResponse": self.allLetterResponse,
            "TargetLetterAvgResponses": self.targetLetterAvgResponses,
            "TargetLetterResponses": self.targetLetterResponses,
            "ColAddRowNum": self.colAddRowNum,
            "ColMulRowNum": self.colMulRowNum,
            "Times": self.times,
            "TotalDataNum": self.totalDataNum,
            "Channels": self.channels,
            "Windows": self.windows,
            "ColAndRowAppear": self.colAndRowAppear
        }

    def __getitem__(self, item: str) -> Any:
        if self._dict.__contains__(item):
            return self._dict[item]
        else:
            raise KeyError("Can't find key {}".format(item))
            # return None

    def __setitem__(self, item: str, data: Any) -> Any:
        if self._dict.__contains__(item):
            self._dict[item] = data

    def LoadData(self) -> None:
        '''
        Description:
            加载数据集
        '''
        self.data = scio.loadmat(self.filePath)
        self.flashing = self.data['Flashing'].astype(np.float64)
        self.signal = self.data['Signal'].astype(np.float64)
        self.stimulusCode = self.data['StimulusCode'].astype(np.float64)
        self.stimulusType = self.data['StimulusType'].astype(np.float64)
        self.targetChar = self.data['TargetChar']
        self.computeEachResponses()
        # self.computeTargetResponses()

    def computeEachResponses(self) -> None:
        '''
        Description:
            获得每个轮次中每行每列的脑电波信号
            同时获得目标字符串的信号
        '''
        for epoch in range(0, np.size(self.signal, 0)):
            m = 0
            times = 0
            rowcolcount = np.zeros(12)
            for n in range(1, np.size(self.signal, 1)):
                if self.flashing[epoch, n] == 0 and self.flashing[epoch, n - 1] == 1:
                    rowcol = int(self.stimulusCode[epoch, n - 1]) - 1
                    self.responses[epoch, rowcol, int(rowcolcount[rowcol]), :, :] = self.signal[epoch,
                                                                                    n - 24:n + self.windows - 24, :]
                    rowcolcount[rowcol] += 1
                    if (int(self.stimulusType[epoch, n - 1]) == 1):
                        # 基于一个想法：每轮只有一个Target，所以每轮只有一个col和row进入上面这个if
                        if (m % 2 == 0):
                            self.targetRowAndCol[epoch, 0] = int(self.stimulusCode[epoch, n - 1])
                            self.targetRowAndCol[epoch, 2] = times
                            self.targetResponses[epoch, 0, times, :, :] = self.signal[epoch,
                                                                          n - 24:n + self.windows - 24, :]
                        else:
                            self.targetRowAndCol[epoch, 1] = int(self.stimulusCode[epoch, n - 1])
                            self.targetResponses[epoch, 1, times, :, :] = self.signal[epoch,
                                                                          n - 24:n + self.windows - 24, :]
                            times += 1
                        m += 1

    def computeTargetResponses(self) -> None:
        '''
        Description:
            获得目标字符串的信号
        '''
        for epoch in range(0, np.size(self.signal, 0)):
            m = 0
            times = 0
            for n in range(1, np.size(self.signal, 1)):
                if self.flashing[epoch, n] == 0 and self.flashing[epoch, n - 1] == 1:
                    # 基于一个想法：每轮只有一个Target，所以每轮只有一个col和row进入上面这个if
                    if (int(self.stimulusType[epoch, n - 1]) == 1):
                        if (m % 2 == 0):
                            self.targetRowAndCol[epoch, 0] = int(self.stimulusCode[epoch, n - 1])
                            self.targetRowAndCol[epoch, 2] = times
                            self.targetResponses[epoch, 0, times, :, :] = self.signal[epoch,
                                                                          n - 24:n + self.windows - 24, :]
                        else:
                            self.targetRowAndCol[epoch, 1] = int(self.stimulusCode[epoch, n - 1])
                            self.targetResponses[epoch, 1, times, :, :] = self.signal[epoch,
                                                                          n - 24:n + self.windows - 24, :]
                            times += 1
                        m += 1

    def avgReps(self) -> np.ndarray:
        '''
        Description:
            对信号的3个维度（行和列出现的次数）取平均
            这会导致第3个维度消失，变成4维向量
        
        Returns:
            meanData:np.ndarray
            返回取平均后的数组
        '''
        return np.mean(self.responses, axis=2)

    # 每个轮次中，每个字符(6*6)的信号
    def getEachLetterAvgResponse(self) -> np.ndarray:
        '''
        Description:
            通过平均行和列的信号获得每个字符的信号
            这里对次数取了平均
            
        Returns:
            allLetterResponse:np.ndarray
            取平均后的数组
        '''
        for epoch in range(np.size(self.signal, 0)):
            aveSignal = np.mean(self.responses[epoch, :, :, :, :], 1)
            m = 0
            for row in range(6, 12):
                for col in range(0, 6):
                    self.allLetterResponse[epoch, :, :] = (aveSignal[row, :, :] + aveSignal[col, :, :]) / 2
                    m += 1
        return self.allLetterResponse

    def getTargetLetterAvgResponse(self) -> np.ndarray:
        '''
        Description:
            获得目标字符的信号
            这里取了平均
            
        Returns:
            targetLetterAvgResponses:np.ndarray
            取平均后的数组
        '''
        for epoch in range(np.size(self.signal), 0):
            aveSignal = np.mean(self.targetResponses[epoch, :, :, :, :], 1)
            self.targetLetterAvgResponses[epoch, :, :] = (aveSignal[0, :, :] + aveSignal[1, :, :]) / 2
        return self.targetLetterAvgResponses

    def getTargetLetterResponse(self) -> np.ndarray:
        '''
        Description:
            获得目标字符的信号
            
        Returns:
            targetLetterResponses:np.ndarray
            目标字符串的信号的数组
        '''
        for epoch in range(np.size(self.targetResponses, 0)):
            for i in range(np.size(self.targetResponses, 2)):
                self.targetLetterResponses[epoch, i, :, :] = (self.targetResponses[epoch, 0, i, :,
                                                              :] + self.targetResponses[epoch, 1, i, :, :]) / 2
        return self.targetLetterResponses

    def alphaTransfNumList(self) -> list:
        '''
        Description:
            目标字符转list
            
        Returns:
            alphaIndex:list
            目标字符列表
        '''
        alphaIndex = [0] * self.times
        m = 0
        for i in self.targetChar[0]:
            num = self.ScreenLetter.index(i)
            alphaIndex[m] = num
            m += 1
        return alphaIndex

    def alphaTransfNdArray(self) -> np.ndarray:
        '''
        Description:
            目标字符转ndarray
        
        Returns:
            alphaNdArray:np.ndarray
            目标字符列表(ndarray类型)
        '''
        alphaNdArray = np.zeros((self.times, len(self.ScreenLetter)))
        m = 0
        for i in self.targetChar[0]:
            num = self.ScreenLetter.index(i)
            alphaNdArray[m][num] = 1
            m += 1
        return alphaNdArray


if __name__ == "__main__":
    TestData = LoadBCICompDataSet(r"E:\文件\Documents\Python\大创\DataSet\BCI_Comp_III_Wads_2004\Subject_A_Train.mat")
    colAddRowNum = TestData['ColAddRowNum']
    windows = TestData['Windows']
    channels = TestData['Channels']
    colMulRowNum = TestData['ColMulRowNum']
    times = TestData['Times']
    Signal = TestData['Signal']
    Flashing = TestData['Flashing']
    StimulusCode = TestData['StimulusCode']
    channel = 10
    StimulusType = TestData['StimulusType']
    TargetChar = TestData['TargetChar']

    responses = np.zeros((colAddRowNum, 15, windows, channels))  # 15是每个信号的行和列出现15次
    letter = np.zeros((colMulRowNum, windows, channels))
    score = np.zeros((colMulRowNum))
    Target = np.zeros((times, windows, channels))
    NonTarget = np.zeros((times, windows, channels))
    Prec = np.zeros(times, dtype='S')
    ScreenLetter = ["A", "B", "C", "D", "E", "F",
                    "G", "H", "I", "J", "K", "L",
                    "M", "N", "O", "P", "Q", "R",
                    "S", "T", "U", "V", "W", "X",
                    "Y", "Z", "1", "2", "3", "4",
                    "5", "6", "7", "8", "9", "_"]

    print(TestData['TargetRowAndCol'])

    # print(Flashing)
    # print(Signal)
    # print(StimulusType)
    # print(StimulusCode)
    # print(TargetChar)

    # print(np.size(Signal,0))

    # log=open("./log.txt",'w')
    # for epoch in range(0,np.size(Signal,0)):
    #     m=0
    #     avgresp=np.mean(TestData['Responses'][epoch,:,:,:,:],axis=1)
    #     for row in range(6,12):
    #         for col in range(0,6):
    #             #计算一个字符的脑电波信息
    #             letter[m,:,:]=(avgresp[row,:,:]+avgresp[col,:,:])/2
    #             #计算一个字符的平均脑电波信号
    #             score[m]=np.mean(letter[m,54-1:124,channel])-np.mean(letter[m,134-1:174,channel])   #一维应该无关轴
    #             m+=1
    #     #print(score)
    #     maxVal=np.max(score)
    #     maxIdx=np.argmax(score)
    #     # print(maxVal)
    #     # print(maxIdx)
    #     Prec[epoch]=ScreenLetter[maxIdx]
    # k=0
    # for i in range(0,np.size(Signal,0)):
    #     if Prec[i].decode('utf-8')==TargetChar[0][i]:
    #         k+=1

    # #print(Prec)
    # correct=(k/np.size(Signal,0))*100;
    # print("correct is {0}%".format(correct))

    # for epoch in range(0,np.size(Signal,0)):
    #     rowcolcount=np.zeros(12)
    #     for n in range(1,np.size(Signal,1)):
    #         if Flashing[epoch,n]==0 and Flashing[epoch,n-1]==1:
    #             rowcol=int(StimulusCode[epoch,n-1])-1
    #             #assert rowcol>=0 and rowcol<12     不报错
    #             responses[rowcol,int(rowcolcount[rowcol]),:,:]=Signal[epoch,n-24:n+windows-24,:]
    #             rowcolcount[rowcol]+=1
    #             #log.write(str(responses.shape))
    #             #log.write("\n")
    #     #下面是matlab中简单的预测器 
    #     m=0
    #     avgresp=np.mean(responses,axis=1)
    #     #print(avgresp.shape)
    #     #avgresp=avgresp.reshape((colAddRowNum,windows,channels))#已经是(colAddRowNum,windows,channels)大小
    #     for row in range(6,12):
    #         for col in range(0,6):
    #             letter[m,:,:]=(avgresp[row,:,:]+avgresp[col,:,:])/2
    #             score[m]=np.mean(letter[m,54-1:124,channel])-np.mean(letter[m,134-1:174,channel])   #一维应该无关轴
    #             m+=1
    #     #print(score)
    #     maxVal=np.max(score)
    #     maxIdx=np.argmax(score)
    #     # print(maxVal)
    #     # print(maxIdx)
    #     Prec[epoch]=ScreenLetter[maxIdx]
    #     label=np.unique(np.multiply(StimulusCode[epoch,:],StimulusType[epoch,:]))
    #     targetLabel=(6*(int(label[2])-7))+int(label[1]);
    #     Target[epoch,:,:]=0.5*(avgresp[int(label[1])-1,:,:]+avgresp[int(label[2])-1,:,:])
    #     NonTarget[epoch,:,:]=np.mean(avgresp,axis=0)-(1/6)*Target[epoch,:,:]

    # k=0
    # for i in range(0,np.size(Signal,0)):
    #     if Prec[i].decode('utf-8')==TargetChar[0][i]:
    #         k+=1

    # #print(Prec)
    # correct=(k/np.size(Signal,0))*100;
    # print("correct is {0}%".format(correct))
    # #print(responses)
    # #log.close()
